import torch
import torch.nn as nn
import torch.nn.functional as F



class AudioClassifier(nn.Module):
    def __init__(self, feature_preprocessor, num_classes=4, in_feats_seq_length=125, in_feats_feature_dim=128):
        super(AudioClassifier, self).__init__()
        self.feature_preprocessor = feature_preprocessor
        self.in_feats_seq_length = in_feats_seq_length
        self.in_feats_feature_dim = in_feats_feature_dim
        self.feature_preprocessor.eval()  # Ensuring the feature extractor does not train
        for param in self.feature_preprocessor.parameters():
            param.requires_grad = False  # Freezing parameters

        # CNN to reduce dimensions spatially
        self.conv_net = nn.Sequential(
            nn.Conv2d(1, 16, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(16, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 8, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )

        # Transformer for capturing temporal dependencies
        transformer_layer = nn.TransformerEncoderLayer(
            d_model=32,  # Matches CNN output channels
            nhead=4,
            dim_feedforward=128,
            dropout=0.1
        )
        self.transformer = nn.TransformerEncoder(transformer_layer, num_layers=1)

        # Classifier to predict classes from pooled features
        cnn_output_dim = (self.in_feats_feature_dim // 8) * 8  # Pooling reduces each spatial dimension by 1/8
        self.inject_layer = nn.Linear(cnn_output_dim, 32)
        self.final_layer = nn.Linear(32, num_classes)

    def forward(self, x):
        with torch.no_grad():
            x = self.feature_preprocessor(x)  # Assume output is [batch_size, seq_length, feature_dim]
        x = x.unsqueeze(1)  # Adding a channel dimension for CNN
        x = self.conv_net(x)  # Applying CNN
        x = x.permute(2, 0, 1, 3)  # Rearranging for transformer input
        x = x.flatten(2)  # Flatten all features for transformer
        x = self.inject_layer(x)
        x = self.transformer(x)  # Transformer processes the sequence
        x = x.mean(dim=0)  # Temporal pooling
        feats = x
        x = self.final_layer(x)  # Final classification
        return x, feats
    
    def get_params(self):
        # Collect trainable parameters excluding the frozen feature extractor
        return [p for p in self.conv_net.parameters() if p.requires_grad] + \
               [p for p in self.transformer.parameters() if p.requires_grad] + \
               [p for p in self.final_layer.parameters() if p.requires_grad]
